#!/usr/bin/python
import os, sys, random
from random import randint
from mako.template import Template
from optparse import OptionParser
import struct, hashlib, socket
TYPE_LIST = ['A']
need_corresponding = ['NS', 'MX', 'CNAME']
def gen_rand_zone(num, zonename='zone'):
    if zonename.endswith('.'):
        zonename = zonename[0:-1]
    name_len = len(str(num))
    return [zonename + '%.*d' % (name_len, i)  + '.' for  i in range(num)]

def create_zone(para, outpath=None):
    makotext="""
% if 'ttl' not in opts or opts['ttl'] == None:
$TTL   100
% else:
$TTL	${opts['ttl']}
% endif
$ORIGIN ${opts['zone']}
@	IN	SOA	ns.${opts['zone']} mail.${opts['zone']} (
			      1		; Serial
			   3600		; Refresh
			  86400		; Retry
			2419200		; Expire
			 604800 )	; Negative Cache TTL
;
@	IN	NS	ns.${opts['zone']}
ns.${opts['zone']}  IN    A       11.22.33.44

"""
    tpl=Template(makotext)
    if para['zone'][-1] != '.':
        para['zone'] += '.'
    txt = tpl.render(opts = para)

    zone_fp = open(para['path'] + '/' + para['zone'] + "txt", "w")
    zone_fp.write(txt)

    qp_file_fp = open(para['qp_file'], 'a')

    if para['rrtype'].upper() == 'A':
        add_a_record(para['domainnum'], para['zone'], para['prefix'], zone_fp, qp_file_fp)
    elif para['rrtype'].upper() == 'TXT':
        add_txt_record(para['domainnum'], para['zone'], para['prefix'], para['service'], zone_fp, qp_file_fp)

    zone_fp.write('\n')
    zone_fp.close()
    qp_file_fp.close()

def gen_txt():
    i = random.randint(0, 1000000000000000)
    return 'http://www.' + hashlib.sha384(str(i)).hexdigest() + '.cnnic.net.cn'

def add_txt_record(num, zone, prefix, type, zone_fp, qp_file_fp):
    domain_len = len(str(num))
    if type.lower() == 'domain':
        for i in xrange(num / 2): 
            domain = prefix + '%.*d' % (domain_len, i) + '.' + zone
            qp_file_fp.write(domain + ' TXT\n')
            domain = domain + '\tIN\tTXT\t' + gen_txt()
            zone_fp.write('\n' + domain) 
            zone_fp.write('\n*.' + domain) 
    else:
        for i in xrange(num / 2): 
            packedip = struct.pack('!I', i)
            domain = socket.inet_ntoa(packedip)
            domain += '.' + zone
            qp_file_fp.write(domain + ' TXT\n')
            domain = domain + '\tIN\tTXT\t' + gen_txt()
            zone_fp.write('\n' + domain) 
            zone_fp.write('\n*.' + domain) 
        
def add_a_record(num, zone, prefix, zone_fp, qp_file_fp):
    domain_len = len(str(num))
    for i in xrange(num):
        domain = prefix + '%.*d' % (domain_len, i) + '.' + zone
        qp_file_fp.write(domain + ' A\n')
        domain = '\n' + domain + '\tIN\tA\t' + gen_rand_ip4()
        zone_fp.write(domain)
def gen_rand_ip4():
    ip4 = str(randint(0, 255))
    ip4 += "."
    ip4 += str(randint(0, 255))
    ip4 += "."
    ip4 += str(randint(0, 255))
    ip4 += "."
    ip4 += str(randint(0, 255))
    return ip4

def main(opts):
    if 'zonenum' not in opts:
        opts['zonenum'] = 1
    if 'zoneprefix' not in opts:
        opts['zoneprefix'] = 'zone'
    if 'domainnum' not in opts:
        opts['domainnum'] = 100
    if 'prefix' not in opts:
        opts['prefix'] = 'www.for.test'
    if 'rrtype' not in opts:
        opts['rrtype'] = 'A'
    if 'service' not in opts:
        opts['service'] = 'domain'

    zonelist = gen_rand_zone(opts['zonenum'], opts['zoneprefix']) 
    for zone in zonelist:
        opts['zone'] = zone
        create_zone(opts, opts['path'])

def getopt():  
    parser = OptionParser()
    parser.add_option('-o', '--path', dest='path')
    parser.add_option('-f', '--file', dest='qp_file')
    parser.add_option('-p', '--prefix', dest='prefix', default = 'www')
    parser.add_option('-n', '--domainnum', dest='domainnum', default=100, type=int)
    parser.add_option('-z', '--zonenum', dest='zonenum', default=1, type=int)
    parser.add_option('-Z', '--zoneprefix', dest='zoneprefix', default = 'zone')
    parser.add_option('-t', '--rrtype', dest='rrtype', default = 'A')
    parser.add_option('-s', '--service', dest='service', default = 'domain')

    opts, args = parser.parse_args()
    if opts.path == None:
       parser.print_help()
       sys.exit(1)
    return opts 

if __name__ == "__main__":
    opts =getopt()
    main(vars(opts))
